#include "task.h"
#include "printk.h"
#include "gate.h"


un_task g_init_task_union __attribute__((__section__(".data.init_task"))) = { TASK_INIT(g_init_task_union.task) };    //链接到Kernel.lds中预留的特殊段空间中
st_mm g_init_mm = {0};
st_thread g_init_thread = {
    .rsp0 = (unsigned long)(g_init_task_union.stack + STACK_SIZE / sizeof(unsigned long)),
    .rsp = (unsigned long)(g_init_task_union.stack + STACK_SIZE / sizeof(unsigned long)),
    .fs = KERNEL_DS,
    .gs = KERNEL_DS,
    .cr2 = 0,
    .trap_nr = 0,
    .error_code = 0
};
st_tss g_init_tss[NR_CPUS] = {
    [0 ... NR_CPUS-1] = TSS_INIT
};

system_call_t g_system_call_table[MAX_SYSTEM_CALL_NR] = {
    [0] = no_system_call,
    [1] = sys_printf,
    [2 ... MAX_SYSTEM_CALL_NR - 1] = no_system_call
};

unsigned long no_system_call(st_pt_regs *regs) {
    color_printk(RED, BLACK, "no_system_call is calling, NR:%#04x\n", regs->rax);
    return -1;
}

unsigned long sys_printf(st_pt_regs *regs) {
    color_printk(BLACK, WHITE, (char *)regs->rdi);
    return 1;
}

void user_level_function() {
    long ret = 0;
    // color_printk(RED, BLACK, "user_level_function task is running\n");
    char string[] = "Hello World!\n";
    __asm__ __volatile__ (
        "leaq sysexit_return_address(%%rip), %%rdx \n\t"
        "movq %%rsp, %%rcx \n\t"
        "sysenter \n\t"
        "sysexit_return_address: \n\t"
        :"=a"(ret)
        :"0"(1), "D"(string)
        :"memory"
    );
    // color_printk(RED, BLACK, "user_level_function task called sysenter, ret:%ld\n", ret);
    while (1);
}

unsigned long do_execve(st_pt_regs *regs) {
    regs->rdx = 0x800000;       // RIP
    regs->rcx = 0xa00000;       // RSP
    regs->rax = 1;
    regs->ds = 0;
    regs->es = 0;
    color_printk(RED, BLACK, "do_execve task is running\n");

    memcpy(user_level_function, (void *)0x800000, 1024);

    return 1;
}

// init进程主函数
unsigned long init(unsigned long arg) {
    st_pt_regs *regs;
    color_printk(RED, BLACK, "init task is running, arg:%#018lx\n", arg);

    CURRENT->thread->rip = (unsigned long)ret_system_call;
    CURRENT->thread->rsp = (unsigned long)CURRENT + STACK_SIZE - sizeof(st_pt_regs);
    regs = (st_pt_regs *)CURRENT->thread->rsp;
    __asm__ __volatile__ (
        "movq %1, %%rsp \n\t"
        "pushq %2 \n\t"
        "jmp do_execve \n\t"
        ::"D"(regs), "m"(CURRENT->thread->rsp), "m"(CURRENT->thread->rip):"memory"
    );
    return 1;
}

inline st_task *get_current() {
    st_task *current = NULL;
    __asm__ __volatile__ ("andq %%rsp, %0 \n\t":"=r"(current):"0"(~32767UL));
    return current;
}

inline void __switch_to(st_task *prev, st_task *next) {
    g_init_tss[0].rsp0 = next->thread->rsp0;
    set_tss64(g_init_tss[0].rsp0, g_init_tss[0].rsp1, g_init_tss[0].rsp2, g_init_tss[0].ist1, g_init_tss[0].ist2, g_init_tss[0].ist3, g_init_tss[0].ist4, g_init_tss[0].ist5, g_init_tss[0].ist6, g_init_tss[0].ist7);

    __asm__ __volatile__ ("movq %%fs, %0 \n\t":"=a"(prev->thread->fs));
	__asm__ __volatile__ ("movq %%gs, %0 \n\t":"=a"(prev->thread->gs));

	__asm__ __volatile__ ("movq	%0, %%fs \n\t"::"a"(next->thread->fs));
	__asm__ __volatile__ ("movq	%0, %%gs \n\t"::"a"(next->thread->gs));

	color_printk(WHITE, BLACK, "prev->thread->rsp0:%#018lx\n", prev->thread->rsp0);
	color_printk(WHITE, BLACK, "next->thread->rsp0:%#018lx\n", next->thread->rsp0);
}

unsigned long do_fork(st_pt_regs *regs, unsigned long clone_flags, unsigned long stack_start, unsigned long stack_size) {
    st_task *tsk = NULL;
    st_thread *thd = NULL;
    st_page *p = NULL;

    color_printk(WHITE, BLACK, "alloc_pages, bitmap:%#018lx\n", *g_mem_mgmt.bits_map);
    p = alloc_pages(ZONE_NORMAL, 1, PG_PTABLE_MAPPED | PG_KERNEL);          // 申请内存空间
    color_printk(WHITE, BLACK, "alloc_pages, bitmap:%#018lx\n", *g_mem_mgmt.bits_map);

    tsk = (st_task *)Phy_To_Virt(p->phy_addr);               // 以申请到的内存作为fork的进程的结构体存放地
    color_printk(WHITE, BLACK, "struct st_task address:%#018lx\n", (unsigned long)tsk);

    memset(tsk, 0, sizeof(st_task));     // 构造新进程的结构信息
    *tsk = *CURRENT;
    list_init(&tsk->list);
    list_add_to_before(&g_init_task_union.task.list, &tsk->list);   // 通过链表链接到第一个进程
    tsk->pid++;
    tsk->state = TASK_UNINTERRUPTIBLE;

    thd = (st_thread *)(tsk + 1);    // 进程结构体后面紧接着存放thread信息
    tsk->thread = thd;

    memcpy(regs, (void *)((unsigned long)tsk + STACK_SIZE - sizeof(st_pt_regs)), sizeof(st_pt_regs)); // 复制栈空间

    thd->rsp0 = (unsigned long)tsk + STACK_SIZE;    // 栈从高地址到低地址使用，低地址存放task信息，共用STACK_SIZE大小
    thd->rip = regs->rip;
    thd->rsp = (unsigned long)tsk + STACK_SIZE - sizeof(st_pt_regs); // 设置栈指针

    if (!(tsk->flags & PF_KTHREAD)) {       // 如果不是系统进程fork，跳转到ret_from_intr执行
        thd->rip = regs->rip = (unsigned long)ret_system_call;
    }
    tsk->state = TASK_RUNNING;
    return 1;
}

unsigned long do_exit(unsigned long code) {
    color_printk(RED, BLACK, "exit task is running, arg:%#018lx\n", code);
    while (1);
}

unsigned long system_call_function(st_pt_regs *regs) {
    return g_system_call_table[regs->rax](regs);
}

extern void kernel_thread_func();
__asm__ (
    "kernel_thread_func: \n\t"
    "   popq %r15 \n\t"
    "   popq %r14 \n\t"
    "	popq %r13 \n\t"	
    "	popq %r12 \n\t"	
    "	popq %r11 \n\t"	
    "	popq %r10 \n\t"	
    "	popq %r9 \n\t"	
    "	popq %r8 \n\t"	
    "	popq %rbx \n\t"	
    "	popq %rcx \n\t"	
    "	popq %rdx \n\t"	
    "	popq %rsi \n\t"	
    "	popq %rdi \n\t"	
    "	popq %rbp \n\t"	
    "	popq %rax \n\t"	
    "	movq %rax, %ds \n\t"
    "	popq %rax \n\t"
    "	movq %rax, %es	\n\t"
    "	popq %rax \n\t"
    "	addq $0x38, %rsp \n\t"
    /////////////////////////////////
    "	movq %rdx, %rdi \n\t"
    "	callq *%rbx \n\t"
    "	movq %rax, %rdi \n\t"
    "	callq do_exit \n\t"
);

int kernel_thread(unsigned long (*fn)(unsigned long), unsigned long arg, unsigned long flags) {
    st_pt_regs regs;
    memset(&regs, 0, sizeof(regs));
    regs.rbx = (unsigned long)fn;
    regs.rdx = (unsigned long)arg;
    regs.ds = KERNEL_DS;
    regs.es = KERNEL_DS;
    regs.cs = KERNEL_CS;
    regs.ss = KERNEL_DS;
    regs.rflags = (1 << 9);
    regs.rip = (unsigned long)kernel_thread_func;

    return do_fork(&regs, flags, 0, 0);
}

void task_init() {
    st_task *p = NULL;
    g_init_mm.pgd = (pml4t_t *)g_global_cr3;
    g_init_mm.start_code = g_mem_mgmt.start_code;
    g_init_mm.end_code = g_mem_mgmt.end_code;
    g_init_mm.start_data = (unsigned long)&_data;
    g_init_mm.end_data = (unsigned long)&_edata;
    g_init_mm.start_rodata = (unsigned long)&_rodata;
    g_init_mm.end_rodata = (unsigned long)&_erodata;
    g_init_mm.start_brk = 0;
    g_init_mm.end_brk = g_mem_mgmt.end_brk;
    g_init_mm.start_stack = _stack_start;

    wrmsr(0x174, KERNEL_CS);    // 设置IA32_SYSTENTER_CS寄存器为64位内核代码段选择子
    color_printk(RED, BLACK, "Debug rsp0:%#018lx, rsp:%#018lx\n", CURRENT->thread->rsp0, CURRENT->thread->rsp);
    wrmsr(0x175, CURRENT->thread->rsp0);
    wrmsr(0x176, (unsigned long)system_call);

    set_tss64(g_init_thread.rsp0, g_init_tss[0].rsp1, g_init_tss[0].rsp2, g_init_tss[0].ist1, g_init_tss[0].ist2, g_init_tss[0].ist3, g_init_tss[0].ist4, g_init_tss[0].ist5, g_init_tss[0].ist6, g_init_tss[0].ist7);

    g_init_tss[0].rsp0 = g_init_thread.rsp0;

    list_init(&g_init_task_union.task.list);

    kernel_thread(init, 10, CLONE_FS | CLONE_FILES | CLONE_SIGNAL);

    g_init_task_union.task.state = TASK_RUNNING;

    p = container_of(list_next(&CURRENT->list), st_task, list);

    SWITCH_TO(CURRENT, p);
}